

from torch.utils.data import Dataset
from typing import Any
import glob 
import os 

from timm.data import create_transform
from torchvision import transforms
from utils.joint_transforms import Compose, RandomHorizontallyFlip, Resize
from PIL import Image
import random 

class MirrorDataset(Dataset):
    def __init__(self, data_dir, is_train=True, image_size=384) -> None:
        super().__init__()
        self.all_images = sorted(glob.glob(f"{data_dir}/*/JPEGImages/*.jpg"))
        self.all_masks = sorted(glob.glob(f"{data_dir}/*/SegmentationClassPNG/*.png"))
        
        print(len(self.all_images), len(self.all_masks))
        
        image_size = image_size
        self.joint_transform = Compose([
            RandomHorizontallyFlip(),
            Resize((image_size, image_size))
        ])
        self.val_joint_transform = Compose([
            Resize((image_size, image_size))
        ])
        self.img_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self.target_transform = transforms.ToTensor()
        self.to_pil = transforms.ToPILImage()
        self.is_train = is_train

    def __getitem__(self, index) -> Any:
        # return super().__getitem__(index)
        image = self.all_images[index]
        mask = self.all_masks[index]
        image = Image.open(image).convert('RGB')
        mask = Image.open(mask).convert('L')

        # transformation
        if self.is_train:
            image, mask = self.joint_transform(image, mask)
        else :
            image, mask = self.val_joint_transform(image, mask)

        image = self.img_transform(image)
       
        mask = self.target_transform(mask)
        
        return {
            "image": image,
            "mask": mask,
        } 
        
    def __len__(self):
        return len(self.all_images)


def get_train_val_dataset(image_size=384):
    pass 
    train_dataset = MirrorDataset("/home/zhaohu/mirror_detection/VMD/train", is_train=True, image_size=image_size)
    val_dataset = MirrorDataset("/home/zhaohu/mirror_detection/VMD/test", is_train=False, image_size=image_size)
    
    return train_dataset, val_dataset

if __name__ == "__main__":
    dataset = MirrorDataset(data_dir="/home/zhaohu/mirror_detection/VMD/train")